# dataset.py
import pandas as pd
from torchvision import transforms
from torch.utils.data import Dataset
from PIL import Image
import os

class FundusDataset(Dataset):
    def __init__(self, csv_path, image_dir, transform=None):
        self.data = pd.read_csv(csv_path)
        self.image_dir = image_dir
        self.transform = transform or transforms.Compose([
            transforms.Resize((224, 224)),
            transforms.RandomHorizontalFlip(),
            transforms.RandomRotation(15),
            transforms.ToTensor(),
            transforms.Normalize([0.485, 0.456, 0.406],
                                 [0.229, 0.224, 0.225])
        ])

    def __len__(self):
        return len(self.data)

    def __getitem__(self, idx):
        row = self.data.iloc[idx]
        img_path = os.path.join(self.image_dir, row['id_code'] + '.png')  # 或 '.jpg'
        image = Image.open(img_path).convert('RGB')
        label = int(row['diagnosis'])
        return self.transform(image), label
